import pymysql


# 数据库连接
def connect_mysql():
    return pymysql.connect(host='localhost', user='root', password='11111111', database='db_student', charset='utf8')


# 执行数据库的增、删、改操作
def query(sql, values):
    with connect_mysql() as conn:
        with conn.cursor() as cursor:
            try:
                cursor.execute(sql, values)
                conn.commit()
                return 1  # 执行成功
            except Exception as e:
                conn.rollback()
                print(f"An error occurred: {e}")
                return 0  # 执行失败


# 单条查询
def select_one(sql: object, *keys: object) -> object:
    conn = connect_mysql()  # 获取数据库连接
    cursor = conn.cursor()  # 获取游标
    try:
        cursor.execute(sql, keys)  # 执行SQL语句
        result = cursor.fetchone()  # 获取单个记录
        return result
    except Exception as e:
        print(f"An error occurred: {e}")
        return None
    finally:
        cursor.close()  # 关闭游标
        conn.close()  # 关闭连接


# 多条查询
def select_all(sql):
    with connect_mysql() as conn:
        with conn.cursor() as cursor:
            try:
                cursor.execute(sql)
                result = cursor.fetchall()
                return result
            except Exception as e:
                print(f"An error occurred: {e}")
                return None


# 测试数据库连接
def test_connection():
    conn = connect_mysql()
    cursor = conn.cursor()
    cursor.execute("SELECT * from information")
    result = cursor.fetchall()
    cursor.close()
    conn.close()
    return result


if __name__ == '__main__':
    result = test_connection()
